package ws

import (
	"fmt"
	"net/http"
	"sync"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

var (
	upgrader    = websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
	connections = make(map[*websocket.Conn]bool)
	connMutex   = sync.Mutex{}
)

// WebsocketHandler WebSocket 服务
func WebsocketHandler(c *gin.Context) {
	// 升级 HTTP 请求为 WebSocket
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		fmt.Println("Upgrade error:", err)
		return
	}

	addConnection(conn)
	fmt.Println("WebSocket client connected!")

	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			fmt.Println("Read error (client disconnected):", err)
			break // 跳出循环表示客户端断开连接
		}
	}

	removeConnection(conn)
	_ = conn.Close()
}

// BroadcastMessage 广播消息到所有 WebSocket 客户端
func BroadcastMessage(message string) {
	connMutex.Lock()
	defer connMutex.Unlock()

	for conn := range connections {
		err := conn.WriteMessage(websocket.TextMessage, []byte(message))
		if err != nil {
			fmt.Println("Write error:", err)
			_ = conn.Close()
			removeConnection(conn)
		}
	}
}

// 添加连接到连接池
func addConnection(conn *websocket.Conn) {
	connMutex.Lock()
	defer connMutex.Unlock()
	connections[conn] = true
}

// 从连接池移除连接
func removeConnection(conn *websocket.Conn) {
	connMutex.Lock()
	defer connMutex.Unlock()
	delete(connections, conn)
}
